#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File Name   : GetResult.py.py
# Author      : moshan
# Mail        : mo_shan@yeah.net
# Version     : 1.0
# Created Time: 2024-03-10 11:08:00
# Function    : 获取SQL语句的表名和字段名
#########################################################################

import sqlparse

import sys, time, re

import json

def f_get_tab(parsed) :

    mark = 0

    global tab

    for i in parsed[0].tokens :

        #碰到子查询就回调一下f_get_tab_manager函数
        if str(i).lower().startswith("(select ") :             
        #if "sqlparse.sql.parenthesis" in str(type(i)).lower() :

            #子查询处理两头括号
            f_get_tab_manager(str(i).lstrip("(").rstrip(")"))  

            mark = 0
            continue

        #mark=1表示前一个关键字是from 及各种join 所以后面的字符串就是表名(别名)
        if mark == 1 and len(str(i)) > 0 and " " != str(i) and "\n" != str(i) and "\t" != str(i) : 
 
            for t in str(i).split(",") : #有可能的join是from a,b where a.id=b.id

                for tmp in t.replace("\n","").replace("\t","").split(" ") :

                    if len(tmp) > 0 : 

                        tab.append(tmp.split(".")[-1].replace("`","")) #如果用db.table就将db.去掉仅留下表名

                        break

            mark = 0

        #判断每次的关键字是否是from join等, 这里的逗号也是一种特殊的join
        tab_mark = str(i).replace(" ","").replace("\n","").replace("\t","").lower()

        if tab_mark in ["innerjoin","leftjoin","rightjoin","join","from"] :

            mark = 1

    return tab

def f_check_operator(s) :

    s = s.replace("\t","")

    res = []

    for tmp in s.split(",") : 

        opt = []

        for i in tmp :

            if i in "+-*/%" : opt.append(i)

        if len(opt) == 0 :

            res.append(tmp)

            continue

        stmp = tmp.replace(" ","")

        for i in opt :

            res.append(stmp.split(i)[0])

            stmp = stmp[stmp.index(i)+1:]

        if len(opt) > 0 :

            res.append(stmp.split(opt[-1])[-1])

        else :

            res.append(stmp)

    return ",".join(res)

def f_edit_str (s,opt) :

    mark = []

    while True :

        if opt not in s : break

        i = s.index(opt)

        s = s[0:i] + s[i+1:]

        mark.append(i)

        if len(mark) == 2 :

            s = s[0:mark[0]] + s[mark[1]+1:]

            mark = []

    return s

def f_check_function(s) :

    fun_list = ["sum","avg","min","max","count","abs","floor","rand","pi","truncate","round","sqrt","insert","upper","left","concat","rtrim","substring","reverse","field","position","substring_index","curdate","current_date","curtime","current_time","now","datediff","adddate","subdate","password","md5","row_number","rank","dense_rank","lag","lead","first_value","last_value","if","ifnull","format","bin","hex","oct","inet_ntoa","inet_aton","cast","json_extract","from_unixtime"]

    res = []

    for tmp in s.split("\n") : 

        tmp = tmp.split(" as ")[0] #把as的别名去掉

        if "(" not in tmp : 

            res.append(tmp)

            continue

        for fun in fun_list :

            stmp = tmp

            if "(" in tmp and "(select " not in tmp.lower() : 

                stmp = tmp.replace(" ","")

            while True :

                if stmp.startswith(fun+"(") :

                    i = stmp.index(fun)+len(fun)+1

                    stmp = stmp[stmp.index(fun)+len(fun)+1:]

                elif ","+fun+"(" in stmp or "("+fun+"(" in stmp : 

                    i = stmp.index(fun)+len(fun)+1

                    stmp = stmp[0:i-len(fun)-1] + stmp[stmp.index(fun)+len(fun)+1:]

                else :

                    break

        stmp = stmp.replace(")","")

        res.append(stmp)

    s = ",".join(res)

    return s

def f_check_function_manager(s,opt) :

    s = s.replace("\t","").lower()

    s = f_edit_str(s,"'")

    s = f_edit_str(s,'"')

    res = []

    for ss in s.split(",") :

        if opt == "func" : res.append(f_check_function(ss))

        else : res.append(f_check_operator(ss))

    s = ",".join(res)

    return s

def f_get_col(parsed) :

    """
    获取sql语句的字段列表, 目前已知的是如果使用了表达式, 聚合函数这种可能会导致解析不全
    """

    mark = 0

    global col

    tmp = []

    for i in parsed[0].tokens : 

        #碰到select关键字就将mark=1, 表明后面就是字段列表
        if str(i).lower() == "select" : mark = 1 

        #字段列表只需要提取identifierlist identifier两个属性, 需要注意这两个属性也可能是代表 表名, 所以需要用上面的逻辑控制一下, 就是说如有没有将mark=1的逻辑, 获取的结果可能是表名列表
        if mark == 1 and ("sqlparse.sql.identifierlist" in str(type(i)).lower() or "sqlparse.sql.identifier" in str(type(i)).lower() or "sqlparse.sql.operation" in str(type(i)).lower() or "sqlparse.sql.function" in str(type(i)).lower()) :

            i = f_check_function_manager(str(i),"func")

            i = f_check_function_manager(str(i),"opt")

            tmp.append(i)

        tab_mark = str(i).replace(" ","").replace("\n","").replace("\t","").lower()

        if tab_mark in ["innerjoin","leftjoin","rightjoin","join","from"] :

            mark = 0

    #print(tmp)

    mark = 0

    tmp_select = []

    for i in parsed[0].tokens : 

        #处理select里面的子查询
        if mark == 0 and str(i).lower().startswith("(select ") :

            tmp_select.append(str(i))

        tab_mark = str(i).replace(" ","").replace("\n","").replace("\t","").lower()

        #from join之后的就不用管了
        if tab_mark in ["innerjoin","leftjoin","rightjoin","join","from"] :

            mark = 1

        elif tab_mark in ["union"] : #如果遇到union就重新计数

            mark = 0

    mark = 0

    if len(tmp_select) > 0 :

        for sql in tmp_select :

            f_get_col_manager(sql.lstrip("(").rstrip(")"))
   
    if len(tmp) < 1 : return 0

    for i in ",".join(tmp).split(",") :

        #print(i,"++++")

        col_mark = str(i).replace(" ","").replace("\n","").replace("\t","").lower()

        if col_mark.startswith("casewhen") :

            #case when这种直接忽略, case col when这种才会关注

            continue

        col_mark = str(i).replace("\n","").replace("\t","").lower()

        if col_mark.startswith("case ") :

            c = col_mark.split("case")[1].split("when")[0].replace(" ","").replace("\n","").replace("\t","").replace("'","").replace('"',"")

        elif " when " in col_mark and (" end " in col_mark or col_mark.endswith(" end")):
            #when 后面的忽略
            continue

        else :

            l = i.strip().split(" ")
    
            c = l[0].replace(" ","").replace("\n","").replace("\t","").replace("'","").replace('"',"").split("(")

            if len(c) > 1 :

                c = c[1].split(")")[0]

            else :

                c = c[0].split(")")[0]

        c = c.split(".")

        if len(c) > 1 :

            c = c[1]

        else :

            c = c[0]

        if c not in col and c.isdigit() is False : 

            if len(c.replace("`","")) > 0 :

                col.append(c.replace("`","")) #去重及纯数字不加入列表

    return col

def f_get_tab_manager(sql) :

    """
    根据sql获取表名列表, 碰到子查询可能会被f_get_tab回调
    """

    parsed = sqlparse.parse(sql)

    return f_get_tab(parsed)

def f_get_col_manager(sql) :

    """
    根据sql获取字段列表, 碰到子查询可能会被f_get_col回调, 没法关联字段跟表名, 所以需要自己处理哪个字段是哪个表的, 如果重复的字段级别应该是一样的, 所以也可以忽略这个问题
    """
    parsed = sqlparse.parse(sql)

    return f_get_col(parsed)

def GetResult(sql):

    code = -1

    d_level = 2

    code = 0

    sql = sql.replace(". ",".")
    
    res["sql"] = sql

    if len(sql) > 500 and "in" in sql.lower() :
    
        sql = re.sub(r'\bin\s*\([^)]*\)', 'IN ()', sql, flags=re.IGNORECASE)
    
    sql = sqlparse.format(sql, reindent = True, keyword_case = 'upper', indent_width = 4)
 
    res["tab"] = f_get_tab_manager(sql)

    res["col"] = f_get_col_manager(sql)
    
    col_list = "('" + "','".join(res["col"]) + "')"

    tab_list = "('" + "','".join(res["tab"]) + "')"

    print(json.dumps(res))

sql_list = []

sql_list.append("select `sql`,id,db_name,db_port,db_type from tb_export_data_new where workorder_id = 1")

sql_list.append("select t1.name,t2.order_id from tb_export_data_new t1 join t2 on t1.id = t2.id where workorder_id = 1")

sql_list.append("select t1.name,t2.order_id,(select t3.addr from t3 limit 1) from tb_export_data_new t1 join t2 on t1.id = t2.id where workorder_id = 1")

sql_list.append("select `sql`,id,db_name_old,db_port,db_type from tb_export_data_old where workorder_id = 1 union  select `sql`,id,db_name,db_port,db_type from tb_export_data_new where workorder_id = 2")

for sql in sql_list :

    tab = []

    col = []

    res = {}

    GetResult(sql)
